import os
import glob
import torch
import random
from PIL import Image
from torch.utils import data
from torchvision import transforms as T
import torchvision.transforms as transforms
class data_prefetcher():
    def __init__(self, loader):
        self.loader = loader
        self.dataiter = iter(loader)
        self.stream = torch.cuda.Stream()
        self.mean = torch.tensor([0.485, 0.456, 0.406]).cuda().view(1, 3, 1, 1)
        self.std = torch.tensor([0.229, 0.224, 0.225]).cuda().view(1, 3, 1, 1)
        # With Amp, it isn't necessary to manually convert data to half.
        # if args.fp16:
        #     self.mean = self.mean.half()
        #     self.std = self.std.half()
        self.num_images = len(loader)
        self.preload()

    def preload(self):
        try:
            self.src_image1, self.src_image2 = next(self.dataiter)
        except StopIteration:
            self.dataiter = iter(self.loader)
            self.src_image1, self.src_image2 = next(self.dataiter)
            
        with torch.cuda.stream(self.stream):
            self.src_image1 = self.src_image1.cuda(non_blocking=True)
            self.src_image1 = self.src_image1.sub_(self.mean).div_(self.std)
            self.src_image2 = self.src_image2.cuda(non_blocking=True)
            self.src_image2 = self.src_image2.sub_(self.mean).div_(self.std)

    def next(self):
        torch.cuda.current_stream().wait_stream(self.stream)
        src_image1 = self.src_image1
        src_image2 = self.src_image2
        self.preload()
        return src_image1, src_image2
    
    def __len__(self):
        """Return the number of images."""
        return self.num_images

class SwappingDataset(data.Dataset):
    """Dataset class for the Artworks dataset and content dataset."""

    def __init__(self, image_dir, img_transform, subffix='jpg', random_seed=1234):
        """Initialize and preprocess the Swapping dataset."""
        self.image_dir = image_dir
        self.img_transform = img_transform
        self.subffix = subffix
        self.dataset = []
        self.random_seed = random_seed
        self.preprocess()
        # self.files = sorted(glob.glob("%s/*.jpg" % self.image_dir))
        self.num_images = len(self.dataset)

    def preprocess(self):
        """Preprocess the Swapping dataset."""
        print("processing Swapping dataset images...")

        temp_path = os.path.join(self.image_dir, '*')
        pathes = glob.glob(temp_path)
        print(temp_path)
        # self.dataset = []
        # 从这里就是错的
        for dir_item in pathes:
            # print("****")
            join_path = glob.glob(os.path.join(dir_item, '*.jpg'))
            temp_list = []
            for item in dir_item:
                temp_list.append(item)
            self.dataset.append(temp_list)
        random.seed(self.random_seed)
        random.shuffle(self.dataset)
        print(len(self.dataset))
        print('Finished preprocessing the Swapping dataset, total dirs number: %d...'% len(self.dataset))
        # print('Finished preprocessing the Swapping dataset, total images number: %d...'% self.num_images)
             
    def __getitem__(self, index):
        """Return two src domain images and two dst domain images."""
        dir_tmp1 = self.dataset[index] # 这句话有问题，self.dataset是空的
        # print(self.dataset)
        dir_tmp1_len = len(self.dataset)
        # filename1 = dir_tmp1[random.randint(0,dir_tmp1_len-1)]
        # filename2 = dir_tmp1[random.randint(0,dir_tmp1_len-1)]
        filename1 = ''.join(self.dataset[random.randint(0, dir_tmp1_len - 1)])
        filename2 = ''.join(self.dataset[random.randint(0, dir_tmp1_len - 1)])
        image1 = self.img_transform(Image.open(filename1))
        image2 = self.img_transform(Image.open(filename2))
        return image1, image2
    
    def __len__(self):
        """Return the number of images."""
        return len(self.dataset)

def GetLoader(dataset_roots, batch_size=16, dataloader_workers=8, random_seed = 1234):
    """Build and return a data loader."""
        
    num_workers = dataloader_workers
    data_root = dataset_roots
    random_seed = random_seed
    
    # c_transforms = []
    #
    # c_transforms.append(T.ToTensor())
    # c_transforms.append(T.Resize(256))
    # c_transforms = T.Compose(c_transforms)

    c_transforms = transforms.Compose([
        # 这里改了改256
        # transforms.Resize((256, 256)),
        # transforms.CenterCrop((256, 256)),
        transforms.ToTensor()
        ])

    content_dataset = SwappingDataset(
                            data_root, 
                            c_transforms,
                            "jpg",
                            random_seed)
    content_data_loader = data.DataLoader(dataset=content_dataset,batch_size=batch_size,
                    drop_last=True, shuffle=False, num_workers=num_workers, pin_memory=True)
    print('content dataset size: ',len(content_dataset))
    prefetcher = data_prefetcher(content_data_loader)
    return prefetcher

def denorm(x):
    out = (x + 1) / 2
    return out.clamp_(0, 1)